import numpy as np
import cv2


def distance_alpha(rgb_img, bg_colors, tolerance, feather, edge_clarity):


    # TODO 代码粘贴区开始
    # 获取尺寸
    height, width = rgb_img.shape[:2]

    # 将参数名适配原算法中的变量
    background_colors = bg_colors  # 与原算法保持一致的局部变量名

    # 使用向量化操作替代双重循环 - 关键优化点！
    # 对图像进行高斯模糊以减少噪点
    smoothed = cv2.GaussianBlur(rgb_img, (3, 3), 0)

    # 将背景颜色转换为 numpy 数组 (N,3)
    bg_colors_array = np.array(background_colors, dtype=np.float32)

    # 重塑图像为 (height*width, 3)
    pixels_reshaped = smoothed.reshape(-1, 3).astype(np.float32)

    # ---------------- 内存优化 ----------------
    # 逐个背景颜色计算距离并取最小值，
    # 避免一次性 broadcast 生成 (像素数 × 颜色数 × 3) 的超大临时数组。
    min_distances = np.full(pixels_reshaped.shape[0], np.inf, dtype=np.float32)
    for col in bg_colors_array:
        diff = pixels_reshaped - col  # (N,3)
        dist = np.sqrt(np.sum(diff * diff, axis=1))  # 欧氏距离
        np.minimum(min_distances, dist, out=min_distances)  # 原地更新最小值

    # 归一化到 0-255
    norm_distances = np.clip(min_distances * (255 / tolerance), 0, 255)

    # 重塑回原图尺寸
    color_distance = norm_distances.reshape(height, width).astype(np.uint8)

    # 将距离图转换为初始掩码
    base_mask = color_distance

    # 对基础掩码应用自适应阈值处理，使边缘更加清晰
    _, thresholded = cv2.threshold(base_mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # 使用形态学操作清理掩码中的噪点
    kernel = np.ones((3, 3), np.uint8)
    cleaned_mask = cv2.morphologyEx(thresholded, cv2.MORPH_OPEN, kernel)
    cleaned_mask = cv2.morphologyEx(cleaned_mask, cv2.MORPH_CLOSE, kernel)

    # 对比度增强以提高边缘清晰度
    alpha_val = float(edge_clarity) * 0.2  # 根据滑块值调整，控制对比度
    enhanced_mask = cv2.convertScaleAbs(cleaned_mask, alpha=alpha_val, beta=0)

    # 应用边缘检测
    edges = cv2.Canny(enhanced_mask, 50, 150)

    # 创建前景和背景区域的确定区域
    sure_fg = cv2.erode(enhanced_mask, kernel, iterations=2)
    sure_bg = cv2.dilate(enhanced_mask, kernel, iterations=2)

    # 将边缘区域与前景区域结合，创建精确的前景掩码
    refined_fg = cv2.bitwise_or(sure_fg, edges)

    # 创建Trimap (确定前景、确定背景和不确定区域)
    trimap = np.zeros((height, width), dtype=np.uint8)
    trimap[sure_bg == 0] = 0  # 背景
    trimap[refined_fg > 0] = 255  # 前景

    uncertain = cv2.bitwise_xor(sure_bg, refined_fg)
    trimap[uncertain > 0] = 128

    # 通过距离变换创建平滑的过渡区
    dist_transform = cv2.distanceTransform(uncertain, cv2.DIST_L2, 3)
    cv2.normalize(dist_transform, dist_transform, 0, 1.0, cv2.NORM_MINMAX)
    alpha_uncertain = (dist_transform * 255).astype(np.uint8)

    # 汇总 alpha
    final_alpha = np.zeros((height, width), dtype=np.uint8)
    final_alpha[trimap == 255] = 255
    final_alpha[trimap == 0] = 0
    final_alpha[trimap == 128] = alpha_uncertain[trimap == 128]

    # 应用边缘羽化处理
    if feather > 0:
        edge_area = cv2.dilate(refined_fg, kernel, iterations=int(feather)) - cv2.erode(refined_fg, kernel, iterations=int(feather))
        edge_alpha = final_alpha.copy()
        blur_size = int(feather * 2) * 2 + 1
        edge_alpha = cv2.GaussianBlur(edge_alpha, (blur_size, blur_size), feather)
        final_alpha[edge_area > 0] = edge_alpha[edge_area > 0]

    # 进一步优化边缘
    final_alpha = cv2.bilateralFilter(final_alpha, 9, 75, 75)

    return final_alpha

    # TODO 代码粘贴区结束
